support for image attachments when classifying questions

Author: escattone

kitsune/llm/questions/classifiers.py

@@ -32,6 +32,7 @@ def classify_question(question: "Question") -> dict[str, Any]:
     payload: dict[str, Any] = {
         "subject": question.title,
         "question": question.content,
+        "image_urls": [image.get_absolute_url() for image in question.get_images()],
         "product": product,
         "topics": get_taxonomy(
             product, include_metadata=["description", "examples"], output_format="JSON"

kitsune/llm/questions/prompt.py

@@ -1,5 +1,7 @@
 from langchain.output_parsers import ResponseSchema, StructuredOutputParser
-from langchain.prompts import ChatPromptTemplate
+from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
+from langchain.schema import HumanMessage
+from langchain.schema.runnable import RunnableLambda
 
 SPAM_INSTRUCTIONS = """
 # Role and goal
@@ -120,17 +122,49 @@
 )
 
 
-spam_prompt = ChatPromptTemplate(
+spam_prompt_template = ChatPromptTemplate(
     (
         ("system", SPAM_INSTRUCTIONS),
-        ("human", USER_QUESTION),
+        MessagesPlaceholder("human_message"),
     )
 ).partial(format_instructions=spam_parser.get_format_instructions())
 
 
-topic_prompt = ChatPromptTemplate(
+topic_prompt_template = ChatPromptTemplate(
     (
         ("system", TOPIC_INSTRUCTIONS),
-        ("human", USER_QUESTION),
+        ("human_message"),
     )
 ).partial(format_instructions=topic_parser.get_format_instructions())
+
+
+def create_human_message(inputs: dict) -> dict:
+    """
+    Creates the human message, with the image URL's if they're present, and
+    then adds it to the inputs dict. Returns the modified inputs dict.
+    """
+    content: list[dict] = [
+        {
+            "type": "text",
+            "text": USER_QUESTION.format(**inputs),
+        },
+    ]
+
+    for image_url in inputs.get("image_urls", ()):
+        content.append(
+            {
+                "type": "image_url",
+                "image_url": {
+                    "url": image_url,
+                },
+            }
+        )
+
+    inputs["human_message"] = [HumanMessage(content=content)]
+    return inputs
+
+
+spam_prompt = RunnableLambda(create_human_message) | spam_prompt_template
+
+
+topic_prompt = RunnableLambda(create_human_message) | topic_prompt_template